﻿#pragma once

#include <iostream>
#include <algorithm>
#include <map>
#include <set>
#include <vector>
#include <cassert>

using namespace std;

template <class T>
void swap(T& a, T& b)
{
	T c(a);
	a = b;
	b = c;
}
template<class K, class V>
struct AVLTreeNode
{
	// 需要 parent 指针，后续更新平衡因⼦可以看到

	pair<K, V> _kv;
	AVLTreeNode<K, V>* _left;
	AVLTreeNode<K, V>* _right;
	AVLTreeNode<K, V>* _parent;
	int _bf; // 平衡因子(balance factor)
	AVLTreeNode(const pair<K, V>& kv)
		:_kv(kv)
		, _left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _bf(0)
	{
	}
};

template<class K, class V>
class AVLTree
{
	typedef AVLTreeNode<K, V> Node;
public:
	Node* GetRoot()
	{
		return _root;
	}
	void InOrder(Node* root)
	{
		if (!root) return;
		InOrder(root->_left);
		cout << root->_kv.first << '-' << root->_kv.second << ' ';
		InOrder(root->_right);
	}
	int Size(Node* node)
	{
		if (node == nullptr) return 0;

		// 当前结点 + 左子树 + 右子树
		return 1 + Size(node->_left) + Size(node->_right);
	}
	void RefreshBF(Node*& parent, Node*& cur)
	{
		while (parent)
		{
			if (cur == parent->_left)       (parent->_bf)--;
			else if (cur == parent->_right) (parent->_bf)++;

			if (parent->_bf == 0)           break;
			else if (abs(parent->_bf) == 2) break;
			else if (abs(parent->_bf) == 1)
			{
				cur = parent;
				parent = parent->_parent;
			}
			else    assert(false);
		}
	}
	void Rotate(Node*& parent, Node*& cur)
	{
		if (parent && parent->_bf == 2)
		{
			if (cur->_bf == 1) RotateL(parent);
			else if (cur->_bf == -1) RotateRL(parent);
		}
		else if (parent && parent->_bf == -2)
		{
			if (cur->_bf == -1) RotateR(parent);
			else if (cur->_bf == 1) RotateLR(parent);
		}
	}
	bool Insert(const pair<K, V>& kv) // 插入函数
	{
		if (_root == nullptr)
		{
			_root = new Node(kv);
			return true;
		}
		Node* parent = nullptr;
		Node* cur = _root;
		while (cur)
		{
			if (cur->_kv.first < kv.first)
			{
				parent = cur;
				cur = cur->_right;
			}
			else if (cur->_kv.first > kv.first)
			{
				parent = cur;
				cur = cur->_left;
			}
			else
			{
				return false;
			}
		}
		cur = new Node(kv);
		if (parent->_kv.first < kv.first)
		{
			parent->_right = cur;
		}
		else
		{
			parent->_left = cur;
		}
		cur->_parent = parent;

		while (parent)
		{
			if (cur == parent->_left) (parent->_bf)--; // 左子树发生改变, 父结点的平衡因子值 - 1
			else if (cur == parent->_right) (parent->_bf)++; // 右子树发生改变, 父结点的平衡因子值 + 1

			if (parent->_bf == 0)      break; // 更新后平衡因子值为 0, 停止更新
			else if (abs(parent->_bf) == 2) break; // 更新后平衡因子值为 2/-2, 需要进行旋转, 同样停止更新
			else if (abs(parent->_bf) == 1)        // 更新后平衡因子值为 1/-1, 继续更新
			{
				cur = parent;
				parent = parent->_parent;
			}
			else    assert(false); // 正常不可能执行到这种情况, 如果遇到了说明前面出错了, 断言报错
		}
		// 判断是否进行旋转操作
		// ...
		if (parent && parent->_bf == 2)
		{
			if (cur->_bf == 1) RotateL(parent);
			else if (cur->_bf == -1) RotateRL(parent);
		}
		else if (parent && parent->_bf == -2)
		{
			if (cur->_bf == -1) RotateR(parent);
			else if (cur->_bf == 1) RotateLR(parent);
		}

		return true;
	}
	void RotateR(Node* parent) // 右单旋
	{
		Node* childL = parent->_left; // 旋转子树的第一个左孩子, 也就是 b 的父亲
		Node* b = childL->_right; // b 结点
		Node* grandparent = parent->_parent; // 待旋转子树根节点的父结点
		// 当从整棵树的根节点开始旋转时, 这个值为 nullptr
		// 移动 b 结点
		parent->_left = b;
		if (b) b->_parent = parent; // b 可能为空, 在访问其成员前需要进行检查

		// 移动首个左孩子结点到转轴节点 parent 的位置
		childL->_right = parent;
		parent->_parent = childL;

		// 更新祖父结点和新根节点的指向
		childL->_parent = grandparent;
		if (grandparent) // 由于祖父结点可能为空, 这时需要进行判断
		{
			// 转轴节点可能位于祖父结点的左/右, 也需要进行额外判断
			if (grandparent->_left == parent) grandparent->_left = childL;
			else                              grandparent->_right = childL;
		}
		else // 这种情况下祖父结点不存在, 也就是说新的根节点 childL 同时也是整棵树的根节点
		{
			_root = childL; // 更新整棵树的根节点 _root
		}

		parent->_bf = childL->_bf = 0; // 新根节点和转轴结点的平衡因子在旋转后一定是 0
		// 并且旋转前后, 平衡因子发生变化的结点只有这两个
		// 不信自己结合实例验证
	}

	// 只需把 RotateR 里的 _left 改成 _right, _right 改成 _left 就行了
	void RotateL(Node* parent) // 左单旋
	{
		Node* childR = parent->_right;
		Node* b = childR->_left;
		Node* grandparent = parent->_parent;

		parent->_right = b;
		if (b) b->_parent = parent;

		childR->_left = parent;
		parent->_parent = childR;

		childR->_parent = grandparent;
		if (grandparent)
		{
			if (grandparent->_right == parent) grandparent->_right = childR;
			else                              grandparent->_left = childR;
		}
		else
		{
			_root = childR;
		}

		parent->_bf = childR->_bf = 0;
	}
	void RotateLR(Node* parent) // 左右双旋
	{
		Node* childL = parent->_left;
		Node* childLR = childL->_right;
		int bf = childLR->_bf;

		RotateL(childL); // 对转轴结点的左孩子左单旋
		RotateR(parent); // 对转轴结点右单旋
		// 只有更新平衡因子的逻辑不一样
		if (bf == 0) // b 初始为空树
		{
			parent->_bf = 0;
			childL->_bf = 0;
			childLR->_bf = 0;
		}
		else if (bf == 1) // b 初始不为空, 新结点插入到 b 的右孩子
		{
			parent->_bf = 0;
			childL->_bf = -1;
			childLR->_bf = 0;
		}
		else if (bf == -1) // b 初始不为空, 新结点插入到 b 的左孩子
		{
			parent->_bf = 1;
			childL->_bf = 0;
			childLR->_bf = 0;
		}
		else // 这种情况正常是不可能出现的, 如果出现了说明前面的逻辑出错了, 就断言报错
		{
			assert(false);
		}
	}
	void RotateRL(Node* parent) // 右左双旋
	{
		Node* childR = parent->_right;
		Node* childRL = childR->_left;
		int bf = childRL->_bf;

		RotateR(childR);
		RotateL(parent);

		if (bf == 0)
		{
			parent->_bf = 0;
			childR->_bf = 0;
			childRL->_bf = 0;
		}
		else if (bf == -1)
		{
			parent->_bf = 0;
			childR->_bf = 1;
			childRL->_bf = 0;
		}
		else if (bf == 1)
		{
			parent->_bf = -1;
			childR->_bf = 0;
			childRL->_bf = 0;
		}
		else
		{
			assert(false);
		}
	}
	Node* Find(const K& key)
	{
		Node* cur = _root;
		while (cur)
		{
			if (cur->_kv.first < key)
			{
				cur = cur->_right;
			}
			else if (cur->_kv.first > key)
			{
				cur = cur->_left;
			}
			else
			{
				return cur;
			}
		}
		return nullptr;
	}
	int Height(Node* root)
	{
		if (root == nullptr) return 0;
		int leftHeight = Height(root->_left);
		int rightHeight = Height(root->_right);
		return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
	}
	bool IsBalanceTree(Node* root)
	{
		// 空树也是AVL树

		if (nullptr == root) return true;

		// 计算 root 结点的平衡因子: 即 root 左右⼦树的高度差

		int leftHeight = Height(root->_left);
		int rightHeight = Height(root->_right);
		int diff = rightHeight - leftHeight;
		// 如果计算出的平衡因子与 root 的平衡因⼦不相等, 或者
		// root 平衡因子的绝对值超过 1, 则一定不是AVL树

		if (abs(diff) >= 2)
		{
			cout << root->_kv.first << "高度差异常" << endl;
			return false;
		}
		if (root->_bf != diff)
		{
			cout << root->_kv.first << "平衡因子异常" << endl;
			return false;
		}
		// root 的左和右如果都是 AVL 树, 则该树一定是AVL树
		return IsBalanceTree(root->_left) && IsBalanceTree(root->_right);
	}
	bool Erase(const pair<K, V>& kv)
	{
		if (!_root)
		{
			return false; // 树为空
		}

		Node* cur = _root;
		Node* pcur = _root;

		// 找到待删除节点
		while (cur)
		{
			if (kv.first > cur->_kv.first)
			{
				pcur = cur;
				cur = cur->_right;
			}
			else if (kv.first < cur->_kv.first)
			{
				pcur = cur;
				cur = cur->_left;
			}
			else
			{
				break; // 找到了要删除的节点
			}
		}

		if (!cur)
		{
			return false; // 没有找到节点
		}

		Node* delParent = cur->_parent; // 保存删除节点的父节点

		if (cur->_left && cur->_right) // 有左右子节点
		{
			Node* pminRight = cur;
			Node* minRight = cur->_right;

			// 寻找右子树的最小节点
			while (minRight->_left)
			{
				pminRight = minRight;
				minRight = minRight->_left;
			}

			if (pminRight->_left == minRight)
			{
				pminRight->_left = minRight->_right;
				if (minRight->_right)
					minRight->_right->_parent = pminRight;
			}
			else if (pminRight->_right == minRight)
			{
				pminRight->_right = minRight->_right;
				if (minRight->_right)
					minRight->_right->_parent = pminRight;
			}

			// 交换节点的值
			std::swap(minRight->_kv, cur->_kv);
			delete minRight;

			// 更新平衡因子的逻辑
			while (delParent)
			{
				if (cur == delParent->_left)
					delParent->_bf--;
				else if (cur == delParent->_right)
					delParent->_bf++;

				// 平衡因子更新完毕后检查是否需要旋转
				if (delParent->_bf == 0) // 平衡因子为0，树已平衡
					break;
				else if (abs(delParent->_bf) == 2) // 需要旋转
					break;

				// 向上回溯
				cur = delParent;
				delParent = delParent->_parent;
			}

			// 执行旋转
			if (delParent && delParent->_bf == 2)
			{
				if (cur->_bf == 1)
					RotateL(delParent);
				else if (cur->_bf == -1)
					RotateRL(delParent);
			}
			else if (delParent && delParent->_bf == -2)
			{
				if (cur->_bf == -1)
					RotateR(delParent);
				else if (cur->_bf == 1)
					RotateLR(delParent);
			}

			return true;
		}

		// 只有一个子节点或没有子节点的情况
		Node* exist = nullptr;
		if (cur->_left && !cur->_right)      exist = cur->_left;
		else if (!cur->_left && cur->_right) exist = cur->_right;

		// 处理根节点情况
		if (cur == _root)
		{
			_root = exist;
			if (_root)
				_root->_parent = nullptr;
		}
		else
		{
			if (kv.first > pcur->_kv.first)
			{
				pcur->_right = exist;
			}
			else if (kv.first < pcur->_kv.first)
			{
				pcur->_left = exist;
			}
			if (exist)
				exist->_parent = pcur;
		}

		delete cur;

		// 从删除节点的父节点开始回溯更新平衡因子
		while (delParent)
		{
			if (cur == delParent->_left)
				delParent->_bf--;
			else if (cur == delParent->_right)
				delParent->_bf++;

			// 平衡因子更新完毕后检查是否需要旋转
			if (delParent->_bf == 0)
				break;
			else if (abs(delParent->_bf) == 2)
				break;

			// 向上回溯
			cur = delParent;
			delParent = delParent->_parent;
		}

		// 执行旋转
		if (delParent && delParent->_bf == 2)
		{
			if (cur->_bf == 1)
				RotateL(delParent);
			else if (cur->_bf == -1)
				RotateRL(delParent);
		}
		else if (delParent && delParent->_bf == -2)
		{
			if (cur->_bf == -1)
				RotateR(delParent);
			else if (cur->_bf == 1)
				RotateLR(delParent);
		}

		return true;
	}


	//bool Erase(const pair<K, V>& kv)
	//{
	//	if (!_root) return false;
	//	Node* cur = _root;
	//	Node* pcur = _root;
	//	while (cur)
	//	{
	//		if (kv.first > cur->_kv.first)
	//		{
	//			pcur = cur;
	//			cur = cur->_right;
	//		}
	//		else if (kv.first < cur->_kv.first)
	//		{
	//			pcur = cur;
	//			cur = cur->_left;
	//		}
	//		else
	//		{
	//			break;
	//		}
	//	}
	//	if (!cur) return false;

	//	Node* delParent = cur->_parent; // 保存要删除结点的父结点, 后面更新循环因子时要用

	//	if (cur->_left && cur->_right)
	//	{
	//		Node* pminRight = cur;
	//		Node* minRight = cur->_right;
	//		while (minRight->_left)
	//		{
	//			pminRight = minRight;
	//			minRight = minRight->_left;
	//		}
	//		if (pminRight->_left == minRight)
	//		{
	//			pminRight->_left = minRight->_right;
	//			if(minRight->_right)
	//			minRight->_right->_parent = pminRight;
	//		}
	//		else if (pminRight->_right == minRight)
	//		{
	//			pminRight->_right = minRight->_right;
	//			if (minRight->_right)
	//			minRight->_right->_parent = pminRight;
	//		}
	//		std::swap(minRight->_kv, cur->_kv); // 交换二者的 kv
	//		delete minRight;

	//		// 更新平衡因子的逻辑, 和 insert 中相差无几
	//		// ---------------------------------------------------------------
	//		while (delParent)
	//		{
	//			if (cur == delParent->_left)
	//				delParent->_bf--;
	//			else if (cur == delParent->_right)
	//				delParent->_bf++;

	//			// 平衡因子更新完毕后检查是否需要旋转
	//			if (delParent->_bf == 0) // 平衡因子为0，树已平衡
	//				break;
	//			else if (abs(delParent->_bf) == 2) // 需要旋转
	//				break;

	//			// 向上回溯
	//			cur = delParent;
	//			delParent = delParent->_parent;
	//		}

	//		// 执行旋转
	//		if (delParent && delParent->_bf == 2)
	//		{
	//			if (cur->_bf == 1)
	//				RotateL(delParent);
	//			else if (cur->_bf == -1)
	//				RotateRL(delParent);
	//		}
	//		else if (delParent && delParent->_bf == -2)
	//		{
	//			if (cur->_bf == -1)
	//				RotateR(delParent);
	//			else if (cur->_bf == 1)
	//				RotateLR(delParent);
	//		}
	//		// ---------------------------------------------------------------

	//		return true;
	//	}

	//	Node* exist = nullptr;
	//	if (cur->_left && !cur->_right)      exist = cur->_left;
	//	else if (!cur->_left && cur->_right) exist = cur->_right;
	//	if (cur == _root)
	//	{
	//		_root = exist;
	//		if (_root) _root->_parent = nullptr;
	//	}
	//	else
	//	{
	//		if (kv.first > pcur->_kv.first)
	//		{
	//			pcur->_right = exist;
	//		}
	//		else if (kv.first < pcur->_kv.first)
	//		{
	//			pcur->_left = exist;
	//		}

	//		if(exist) exist->_parent = pcur;
	//	}
	//	RefreshBF(delParent, cur);
	//	Rotate(delParent, cur);
	//	delete cur;

	//	// ---------------------------------------------------------------
	//	while (delParent)
	//	{
	//		if (cur == delParent->_left)
	//			delParent->_bf--;
	//		else if (cur == delParent->_right)
	//			delParent->_bf++;

	//		// 平衡因子更新完毕后检查是否需要旋转
	//		if (delParent->_bf == 0)
	//			break;
	//		else if (abs(delParent->_bf) == 2)
	//			break;

	//		// 向上回溯
	//		cur = delParent;
	//		delParent = delParent->_parent;
	//	}

	//	// 执行旋转
	//	if (delParent && delParent->_bf == 2)
	//	{
	//		if (cur->_bf == 1)
	//			RotateL(delParent);
	//		else if (cur->_bf == -1)
	//			RotateRL(delParent);
	//	}
	//	else if (delParent && delParent->_bf == -2)
	//	{
	//		if (cur->_bf == -1)
	//			RotateR(delParent);
	//		else if (cur->_bf == 1)
	//			RotateLR(delParent);
	//	}
	//	// ---------------------------------------------------------------


	//	return true;
	//}
private:
	Node* _root = nullptr;
};